import matplotlib.pyplot as plt
import numpy as np
import torch
from matplotlib.colors import Normalize
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms

from Unet_train import UNet  # 确保这个路径和模块名正确


class MedicalImageTestDataset(Dataset):
    def __init__(self):
        self.img_path = '/Volumes/For_Mac/dateset/ces/case0040_slice098.npy'
        self.label_path = '/Volumes/For_Mac/dateset/Synapse_npy/train/labels/case0040_slice098.npy'

    def __len__(self):
        return 1

    def __getitem__(self, idx):
        image = np.load(self.img_path)
        label = np.load(self.label_path)
        transform = transforms.Compose([
            transforms.ToTensor(),
        ])
        image = transform(image)
        label = transform(label).long()
        return image, label


def test(model, dataloader, device):
    model.eval()
    with torch.no_grad():
        for images, labels in dataloader:
            images = images.to(device)
            outputs = model(images)
            # 使用softmax并通过torch.argmax获取最可能的类别
            preds = torch.argmax(outputs, dim=1)
            yield images.cpu(), preds.cpu(), labels.cpu()


if __name__ == '__main__':
    dataset = MedicalImageTestDataset()
    dataloader = DataLoader(dataset, batch_size=1, shuffle=False)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = UNet(in_channels=1, out_channels=9)
    model.load_state_dict(torch.load('../pt_file/synapse-unet-29.pt', map_location=device, weights_only=True))
    model.to(device)

    norm = Normalize(vmin=0, vmax=8)  # 假设预测的类别标签从0到8

    for images, preds, labels in test(model, dataloader, device):
        fig, ax = plt.subplots(1, 3, figsize=(18, 6))
        ax[0].imshow(images[0][0], cmap='gray')
        ax[1].imshow(preds[0], cmap='tab20', norm=norm)
        ax[2].imshow(labels[0][0], cmap='tab20', norm=norm)

        ax[0].title.set_text('Original Image')
        ax[1].title.set_text('Predicted Mask')
        ax[2].title.set_text('True Mask')

        plt.setp(ax, xticks=[], yticks=[])
        plt.show()
